from .base_strategy import BaseStrategy
import backtrader as bt
###
# 使用sma20 的最低价和 sma20 的最高价2根均线形成一个通道
# 做空的逻辑
# 当连续的2根k都处于通道之下时，当前k线实体有触及到通道，并且收盘价处于通道下方
# 这就视为趋势的回调，然后以最近3根k线的最低价，视为关键价格，
# 在未来6根k线内，价格向下突破关键价格，则进行开单
# 使用atr 进行止盈止损，做多的逻辑则刚好相反
# #
class MyStrategy(BaseStrategy):
    params = (
      ('sma_period', 20),       # SMA周期
        ('lookback_bars', 3),     # 回看多少根K线确定关键价格
        ('max_bars', 6),          # 未来多少根K线内有效
        ('atr_period', 14),       # ATR周期
        ('atr_multiple', 2.0),    # ATR倍数用于止盈止损
    )
 
    def __init__(self):
        super().__init__()

        self.init_indicators()

        # 用于跟踪订单
        self.short_signal = False
        self.long_signal = False
        self.key_price = None
        self.signal_bar = 0
        self.order = None
    
    def init_indicators(self):
        """在此定义指标 (EMA/MACD/RSI等)"""
        # SMA20计算
        # self.sma = bt.indicators.SimpleMovingAverage(
        #     self.data.close, period=self.p.sma_period)
        
        # 计算SMA20期间的最高价和最低价，形成通道
        self.sma_high = bt.indicators.SimpleMovingAverage(
            self.data.high, period=self.p.sma_period)
        self.sma_low = bt.indicators.SimpleMovingAverage(
            self.data.low, period=self.p.sma_period)
        
        # ATR用于止盈止损
        self.atr = bt.indicators.ATR(
            self.data, period=self.p.atr_period)
        


    def notify_order(self, order):
        super().notify_order(order)

        if order.status in [order.Submitted, order.Accepted]:
            # 订单已提交/接受 - 不做任何操作
            return
        
        # 重置订单
        self.order = None

    def next(self):
        super().next()
 
        # 检查是否有持仓
        if self.order:
            return
            
        # 检查是否有持仓
        if not self.position:
            self.check_signals()
            self.check_entry()
        else:
            self.check_exit()

    def check_signals(self):
        # 重置信号
        # self.short_signal = False
        # self.long_signal = False
        
        # 做空信号条件
        # 1. 连续2根K线处于通道之下
        # 2. 当前K线实体触及通道
        # 3. 收盘价处于通道下方
        if (self.data.close[-1] < self.sma_low[-1] and 
            self.data.close[-2] < self.sma_low[-2]):

            if self.data.high[0] >= self.sma_low[0] and self.data.close[0] < self.sma_low[0]:
                self.short_signal = True
                # 记录关键价格  并且开始计算k线是否突破
                self.key_price = min(self.data.low[i] for i in range(-self.p.lookback_bars, 1))
                self.signal_bar = len(self) 
        
        # 做多信号条件（与做空相反）
        if (self.data.close[-1] > self.sma_high[-1] and 
            self.data.close[-2] > self.sma_high[-2]):
            if self.data.low[0] <= self.sma_high[0] and self.data.close[0] > self.sma_high[0]:
                self.long_signal = True
                self.key_price = max(self.data.high[i] for i in range(-self.p.lookback_bars, 1))
                self.signal_bar = len(self) 

    def check_entry(self):
        # 检查做空入场条件
        if self.short_signal:
            # 在未来6根K线内价格向下突破关键价格
            if (len(self) <= self.signal_bar + self.p.max_bars ):
                if self.data.low[0] < self.key_price:
                    # 计算止盈止损
                    stop_loss = self.data.close[0] + self.atr[0] * self.p.atr_multiple
                    take_profit = self.data.close[0] - self.atr[0] * self.p.atr_multiple
                    # 发送做空订单
                    self.order = self.sell()
                    # 保存止损止盈价格
                    self.stop_price = stop_loss
                    self.profit_price = take_profit
                    self.short_signal = False
            else:
                self.short_signal = False
        
        # 检查做多入场条件
        elif self.long_signal:
            if (len(self) <= self.signal_bar + self.p.max_bars ):
                if self.data.high[0] > self.key_price:
                    # 计算止盈止损
                    stop_loss = self.data.close[0] - self.atr[0] * self.p.atr_multiple
                    take_profit = self.data.close[0] + self.atr[0] * self.p.atr_multiple
                    # 发送做多订单
                    self.order = self.buy()
                    # 保存止损止盈价格
                    self.stop_price = stop_loss
                    self.profit_price = take_profit
                    self.long_signal = False
            else:
                self.long_signal = False

    def check_exit(self):
        # 检查止损或止盈条件
        if self.position.size > 0:  # 多头仓位
            if self.data.high[0] >= self.profit_price or self.data.low[0] <= self.stop_price:
                self.close()
        elif self.position.size < 0:  # 空头仓位
            if self.data.low[0] <= self.profit_price or self.data.high[0] >= self.stop_price:
                self.close()
